package com.yanggu.flink.datastream_api.function

import com.yanggu.flink.datastream_api.pojo.SensorReading
import org.apache.flink.api.common.accumulators.IntCounter
import org.apache.flink.api.common.functions.RichMapFunction
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.scala._


/**
 * 测试自定义的RichFunction函数
 * open()方法、close()方法生命周期方法在每个并行度上只会被调用一次
 * getRuntimeContext()方法获取运行时上下文例如函 数执行的并行度, 任务的名字
 */
object UDFFunctionDemo {

  def main(args: Array[String]): Unit = {
    val environment = StreamExecutionEnvironment.getExecutionEnvironment

    val stream = environment.fromCollection(List(
      SensorReading("sensor_1", System.currentTimeMillis(), 35.8),
      SensorReading("sensor_2", System.currentTimeMillis(), 45.8),
      SensorReading("sensor_3", System.currentTimeMillis(), 65.8)
    ))

    stream
      //这里设置了map的并行度为2
      .map(new MyMapFunction).setParallelism(2)
      .print("测试自定义RichFunction").setParallelism(1)

    //可以获取执行的结果
    val result = environment.execute("测试自定义RichFunction")

    //这里可以获取累加器的计算结果
    val value = result.getAccumulatorResult[Int]("MyFunction的计数器")
    println(s"MyFunction的计数器计算结果: $value")
  }

}

/**
 * Rich Function是Function函数的Rich版本,
 * 拥有一些声明周期方法, 可以获取运行时的上下文, 因此可以实现一些更复杂的功能
 *
 * @tparam IN  SensorReading 输入的数据类型
 * @tparam Out SensorReading 输出的数据类型
 */
class MyMapFunction extends RichMapFunction[SensorReading, SensorReading] {

  var subtask = 0

  var intCounter: IntCounter = _

  /**
   * open()方法是rich function的初始化方法,
   * 当一个算子例如map或者filter被调用之前open()会被执行, 且只会执行一次
   * 注意这句话
   * 在每个并行度上只会被调用一次
   *
   * @param parameters
   */
  override def open(parameters: Configuration): Unit = {
    //这里可以进行一些初始化的工作, 例如建立一个和hdfs的链接
    //在这里可以获取上下文信息
    this.subtask = getRuntimeContext.getIndexOfThisSubtask
    this.intCounter = getRuntimeContext.getIntCounter("MyFunction的计数器")
    //只会被打印2次, 设置了map的并行度为2
    println(Thread.currentThread() + "open方法被执行了")
  }

  /**
   * 默认生命周期方法, 最后一个方法, 做一些清理工作
   * 注意这句话
   * 在每个并行度上只会被调用一次, 而且是最后调用
   */
  override def close(): Unit = {
    //只会被打印2次, 设置了map的并行度为2
    println(Thread.currentThread() + "close方法被执行了")
    println(Thread.currentThread() + intCounter.toString)
  }

  /**
   * 需要重写的核心方法
   *
   * @param value
   * @return
   */
  override def map(value: SensorReading): SensorReading = {
    println(Thread.currentThread() + "map方法被执行了")
    intCounter.add(1)
    value
  }

}
